from typing import Generic, TypeVar, Union, Dict, List, AsyncIterator

from bson import ObjectId
from pymongo import InsertOne, UpdateOne, UpdateMany, ReplaceOne, DeleteOne, DeleteMany
from typing_extensions import TypedDict
from motor.core import AgnosticCollection

from ..utils.singleton import Singleton


T = TypeVar('T')
OperationType = Union[InsertOne, UpdateOne, UpdateMany, ReplaceOne, DeleteOne, DeleteMany]


class BaseDao(Generic[T], Singleton):

    def get_db(self) -> AgnosticCollection:
        raise NotImplemented('必须实现get_db方法')

    async def insert_one(self, insert_data: T):
        db = self.get_db()
        return await db.insert_one(insert_data)

    async def insert_many(self, data: List[T]):
        db = self.get_db()
        return await db.insert_many(data)

    async def insert_or_update_by_id(self, id_: Union[str, ObjectId], data: T):
        """根据id来更新或者插入一条新的数据"""
        db = self.get_db()
        if isinstance(id_, str):
            id_ = ObjectId(id_)
        query = {
            '_id': id_
        }
        return await db.replace_one(query, data, upsert=True)

    async def insert_or_update_by_query(self, query: Dict, data: T):
        """根据id来更新或者插入一条新的数据"""
        db = self.get_db()
        return await db.replace_one(query, data, upsert=True)

    async def update_by_id(self, id_: Union[str, ObjectId], update_data: Dict):
        db = self.get_db()
        if isinstance(id_, str):
            id_ = ObjectId(id_)
        query = {
            '_id': id_
        }

        update_ = {
            '$set': update_data
        }
        return await db.update_one(query, update_)

    async def update_many_by_query(self, query: Dict, update_data: Dict):
        db = self.get_db()
        update_ = {
            '$set': update_data
        }
        return await db.update_many(query, update_)

    async def get_by_id(self, id_: Union[str, ObjectId]) -> T:
        db = self.get_db()
        if isinstance(id_, str):
            id_ = ObjectId(id_)
        query = {
            '_id': id_
        }
        return await db.find_one(query)

    async def get_by_query(self, query: Dict) -> AsyncIterator[T]:
        db = self.get_db()
        return db.find(query)

    async def delete_by_id(self, id_: Union[str, ObjectId]):
        db = self.get_db()
        if isinstance(id_, str):
            id_ = ObjectId(id_)
        query = {
            '_id': id_
        }
        return await db.delete_one(query)

    async def delete_many_by_query(self, query: Dict):
        db = self.get_db()
        return await db.delete_many(query)

    async def bulk_write(self, data: List[OperationType]):
        db = self.get_db()
        return await db.bulk_write(data)


if __name__ == '__main__':
    class TestEntity(TypedDict):
        abc: str
        ddd: int

    class TestDao(BaseDao[TestEntity]):
        pass

    async def test_method():
        dao = TestDao()
        a = await dao.get_by_id('1')


